// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/builder.h>

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <turbo/base/checked_cast.h>
#include <nebula/util/hashing.h>
#include <nebula/core/visit_type_inline.h>

namespace nebula {

    class MemoryPool;

    // ----------------------------------------------------------------------
    // Helper functions


    // Generic int builder that delegates to the builder for a specific
    // type. Used to reduce the number of template instantiations in the
    // exact_index_type case below, to reduce build time and memory usage.
    class TURBO_EXPORT TypeErasedIntBuilder : public ArrayBuilder {
    public:
        explicit TypeErasedIntBuilder(MemoryPool *pool = default_memory_pool(),
                                      int64_t alignment = kDefaultBufferAlignment)
                : ArrayBuilder(pool, alignment) {
            // Not intended to be used, but adding this is easier than adding a bunch of enable_if
            // magic to builder_dict.h
                    DKCHECK(false);
        }

        explicit TypeErasedIntBuilder(const std::shared_ptr<DataType> &type,
                                      MemoryPool *pool = default_memory_pool(),
                                      int64_t alignment = kDefaultBufferAlignment)
                : ArrayBuilder(pool), type_id_(type->id()) {
                    DKCHECK(is_integer(type_id_));
            switch (type_id_) {
                case Type::UINT8:
                    builder_ = std::make_unique<UInt8Builder>(pool);
                    break;
                case Type::INT8:
                    builder_ = std::make_unique<Int8Builder>(pool);
                    break;
                case Type::UINT16:
                    builder_ = std::make_unique<UInt16Builder>(pool);
                    break;
                case Type::INT16:
                    builder_ = std::make_unique<Int16Builder>(pool);
                    break;
                case Type::UINT32:
                    builder_ = std::make_unique<UInt32Builder>(pool);
                    break;
                case Type::INT32:
                    builder_ = std::make_unique<Int32Builder>(pool);
                    break;
                case Type::UINT64:
                    builder_ = std::make_unique<UInt64Builder>(pool);
                    break;
                case Type::INT64:
                    builder_ = std::make_unique<Int64Builder>(pool);
                    break;
                default:
                            DKCHECK(false);
            }
        }

        void reset() override { return builder_->reset(); }

        turbo::Status append(int32_t value) {
            switch (type_id_) {
                case Type::UINT8:
                    return turbo::checked_cast<UInt8Builder *>(builder_.get())->append(value);
                case Type::INT8:
                    return turbo::checked_cast<Int8Builder *>(builder_.get())->append(value);
                case Type::UINT16:
                    return turbo::checked_cast<UInt16Builder *>(builder_.get())->append(value);
                case Type::INT16:
                    return turbo::checked_cast<Int16Builder *>(builder_.get())->append(value);
                case Type::UINT32:
                    return turbo::checked_cast<UInt32Builder *>(builder_.get())->append(value);
                case Type::INT32:
                    return turbo::checked_cast<Int32Builder *>(builder_.get())->append(value);
                case Type::UINT64:
                    return turbo::checked_cast<UInt64Builder *>(builder_.get())->append(value);
                case Type::INT64:
                    return turbo::checked_cast<Int64Builder *>(builder_.get())->append(value);
                default:
                            DKCHECK(false);
            }
            return turbo::unimplemented_error("Internal implementation error");
        }

        turbo::Status append_null() override { return builder_->append_null(); }

        turbo::Status append_nulls(int64_t length) override { return builder_->append_nulls(length); }

        turbo::Status append_empty_value() override { return builder_->append_empty_value(); }

        turbo::Status append_empty_values(int64_t length) override {
            return builder_->append_empty_values(length);
        }

        turbo::Status append_scalar(const Scalar &scalar, int64_t n_repeats) override {
            return builder_->append_scalar(scalar, n_repeats);
        }

        turbo::Status append_scalars(const ScalarVector &scalars) override {
            return builder_->append_scalars(scalars);
        }

        turbo::Status append_array_slice(const ArraySpan &array, int64_t offset,
                                       int64_t length) override {
            return builder_->append_array_slice(array, offset, length);
        }

        turbo::Status finish_internal(std::shared_ptr<ArrayData> *out) override {
            return builder_->finish_internal(out);
        }

        std::shared_ptr<DataType> type() const override { return builder_->type(); }

    private:
        std::unique_ptr<ArrayBuilder> builder_;
        Type::type type_id_;
    };

    struct DictionaryBuilderCase {
        template<typename VT, typename Enable = typename VT::c_type>
        turbo::Status Visit(const VT &) {
            return CreateFor<VT>();
        }

        turbo::Status Visit(const NullType &) { return CreateFor<NullType>(); }

        turbo::Status Visit(const BinaryType &) { return CreateFor<BinaryType>(); }

        turbo::Status Visit(const StringType &) { return CreateFor<StringType>(); }

        turbo::Status Visit(const LargeBinaryType &) { return CreateFor<LargeBinaryType>(); }

        turbo::Status Visit(const LargeStringType &) { return CreateFor<LargeStringType>(); }

        turbo::Status Visit(const BinaryViewType &) { return CreateFor<BinaryViewType>(); }

        turbo::Status Visit(const StringViewType &) { return CreateFor<StringViewType>(); }

        turbo::Status Visit(const FixedSizeBinaryType &) { return CreateFor<FixedSizeBinaryType>(); }

        turbo::Status Visit(const Decimal128Type &) { return CreateFor<Decimal128Type>(); }

        turbo::Status Visit(const Decimal256Type &) { return CreateFor<Decimal256Type>(); }

        turbo::Status Visit(const DataType &value_type) { return NotImplemented(value_type); }

        turbo::Status Visit(const Fp16Type &value_type) { return NotImplemented(value_type); }

        turbo::Status NotImplemented(const DataType &value_type) {
            return turbo::unimplemented_error(
                    "MakeBuilder: cannot construct builder for dictionaries with value type ",
                    value_type);
        }

        template<typename VT>
        turbo::Status CreateFor() {
            using AdaptiveBuilderType = DictionaryBuilder<VT>;
            if (dictionary != nullptr) {
                out->reset(new AdaptiveBuilderType(dictionary, pool));
            } else if (exact_index_type) {
                if (!is_integer(index_type->id())) {
                    return turbo::failed_precondition_error("MakeBuilder: invalid index type ", *index_type);
                }
                out->reset(new internal::DictionaryBuilderBase<TypeErasedIntBuilder, VT>(
                        index_type, value_type, pool));
            } else {
                auto start_int_size = index_type->byte_width();
                out->reset(new AdaptiveBuilderType(start_int_size, value_type, pool));
            }
            return turbo::OkStatus();
        }

        turbo::Status create() { return visit_type_inline(*value_type, this); }

        MemoryPool *pool;
        const std::shared_ptr<DataType> &index_type;
        const std::shared_ptr<DataType> &value_type;
        const std::shared_ptr<Array> &dictionary;
        bool exact_index_type;
        std::unique_ptr<ArrayBuilder> *out;
    };

    struct MakeBuilderImpl {
        template<typename T>
        enable_if_not_nested<T, turbo::Status> Visit(const T &t) {
            out.reset(new typename TypeTraits<T>::BuilderType(type, pool));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const DictionaryType &dict_type) {
            DictionaryBuilderCase visitor = {pool,
                                             dict_type.index_type(),
                                             dict_type.get_value_type(),
                    /*dictionary=*/nullptr,
                                             exact_index_type,
                                             &out};
            return visitor.create();
        }

        turbo::Status Visit(const ListType &list_type) {
            std::shared_ptr<DataType> value_type = list_type.get_value_type();
            TURBO_MOVE_OR_RAISE(auto value_builder, ChildBuilder(value_type));
            out.reset(new ListBuilder(pool, std::move(value_builder), type));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const LargeListType &list_type) {
            std::shared_ptr<DataType> value_type = list_type.get_value_type();
            TURBO_MOVE_OR_RAISE(auto value_builder, ChildBuilder(value_type));
            out.reset(new LargeListBuilder(pool, std::move(value_builder), type));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const ListViewType &list_view_type) {
            std::shared_ptr<DataType> value_type = list_view_type.get_value_type();
            TURBO_MOVE_OR_RAISE(auto value_builder, ChildBuilder(value_type));
            out.reset(new ListViewBuilder(pool, std::move(value_builder), std::move(type)));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const LargeListViewType &large_list_view_type) {
            std::shared_ptr<DataType> value_type = large_list_view_type.get_value_type();
            TURBO_MOVE_OR_RAISE(auto value_builder, ChildBuilder(value_type));
            out.reset(new LargeListViewBuilder(pool, std::move(value_builder), std::move(type)));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const MapType &map_type) {
            TURBO_MOVE_OR_RAISE(auto key_builder, ChildBuilder(map_type.key_type()));
            TURBO_MOVE_OR_RAISE(auto item_builder, ChildBuilder(map_type.item_type()));
            out.reset(
                    new MapBuilder(pool, std::move(key_builder), std::move(item_builder), type));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const FixedSizeListType &list_type) {
            auto value_type = list_type.get_value_type();
            TURBO_MOVE_OR_RAISE(auto value_builder, ChildBuilder(value_type));
            out.reset(new FixedSizeListBuilder(pool, std::move(value_builder), type));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const StructType &struct_type) {
            TURBO_MOVE_OR_RAISE(auto field_builders, FieldBuilders(*type, pool));
            out.reset(new StructBuilder(type, pool, std::move(field_builders)));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const SparseUnionType &) {
            TURBO_MOVE_OR_RAISE(auto field_builders, FieldBuilders(*type, pool));
            out.reset(new SparseUnionBuilder(pool, std::move(field_builders), type));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const DenseUnionType &) {
            TURBO_MOVE_OR_RAISE(auto field_builders, FieldBuilders(*type, pool));
            out.reset(new DenseUnionBuilder(pool, std::move(field_builders), type));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const RunEndEncodedType &ree_type) {
            TURBO_MOVE_OR_RAISE(auto run_end_builder, ChildBuilder(ree_type.run_end_type()));
            TURBO_MOVE_OR_RAISE(auto value_builder, ChildBuilder(ree_type.get_value_type()));
            out.reset(new RunEndEncodedBuilder(pool, std::move(run_end_builder),
                                               std::move(value_builder), type));
            return turbo::OkStatus();
        }

        turbo::Status Visit(const ExtensionType &) { return NotImplemented(); }

        turbo::Status Visit(const DataType &) { return NotImplemented(); }

        turbo::Status NotImplemented() {
            return turbo::unimplemented_error("MakeBuilder: cannot construct builder for type ",
                                                 type->to_string());
        }

        turbo::Result<std::unique_ptr<ArrayBuilder>> ChildBuilder(
                const std::shared_ptr<DataType> &type) {
            MakeBuilderImpl impl{pool, type, exact_index_type, /*out=*/nullptr};
            TURBO_RETURN_NOT_OK(visit_type_inline(*type, &impl));
            return std::move(impl.out);
        }

        turbo::Result<std::vector<std::shared_ptr<ArrayBuilder>>> FieldBuilders(const DataType &type,
                                                                                MemoryPool *pool) {
            std::vector<std::shared_ptr<ArrayBuilder>> field_builders;
            for (const auto &field: type.fields()) {
                std::unique_ptr<ArrayBuilder> builder;
                MakeBuilderImpl impl{pool, field->type(), exact_index_type, /*out=*/nullptr};
                TURBO_RETURN_NOT_OK(visit_type_inline(*field->type(), &impl));
                field_builders.emplace_back(std::move(impl.out));
            }
            return field_builders;
        }

        MemoryPool *pool;
        const std::shared_ptr<DataType> &type;
        bool exact_index_type;
        std::unique_ptr<ArrayBuilder> out;
    };

    turbo::Status MakeBuilder(MemoryPool *pool, const std::shared_ptr<DataType> &type,
                              std::unique_ptr<ArrayBuilder> *out) {
        MakeBuilderImpl impl{pool, type, /*exact_index_type=*/false, /*out=*/nullptr};
        TURBO_RETURN_NOT_OK(visit_type_inline(*type, &impl));
        *out = std::move(impl.out);
        return turbo::OkStatus();
    }

    turbo::Status MakeBuilderExactIndex(MemoryPool *pool, const std::shared_ptr<DataType> &type,
                                        std::unique_ptr<ArrayBuilder> *out) {
        MakeBuilderImpl impl{pool, type, /*exact_index_type=*/true, /*out=*/nullptr};
        TURBO_RETURN_NOT_OK(visit_type_inline(*type, &impl));
        *out = std::move(impl.out);
        return turbo::OkStatus();
    }

    turbo::Status MakeDictionaryBuilder(MemoryPool *pool, const std::shared_ptr<DataType> &type,
                                        const std::shared_ptr<Array> &dictionary,
                                        std::unique_ptr<ArrayBuilder> *out) {
        const auto &dict_type = static_cast<const DictionaryType &>(*type);
        DictionaryBuilderCase visitor = {
                pool, dict_type.index_type(), dict_type.get_value_type(),
                dictionary, /*exact_index_type=*/false, out};
        return visitor.create();
    }

}  // namespace nebula
